import cv2
import numpy as np
import os
import csv
from datetime import datetime

# 初始化人脸检测器和识别器
face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
recognizer = cv2.face.LBPHFaceRecognizer_create()

# 全局变量
known_face_names = []
face_samples = []
ids = []
id_count = 0

# 创建必要文件夹
if not os.path.exists("faces_data"):
    os.makedirs("faces_data")


# 加载已知人脸数据
def load_known_faces():
    global recognizer, known_face_names, id_count

    if os.path.exists("faces_data/faces.yml"):
        recognizer.read("faces_data/faces.yml")

    if os.path.exists("faces_data/names.csv"):
        with open("faces_data/names.csv", "r") as f:
            known_face_names = [line.strip() for line in f.readlines()]
            id_count = len(known_face_names)

    print(f"已加载 {len(known_face_names)} 个已知人脸")


# 保存人脸数据
def save_face_data():
    recognizer.write("faces_data/faces.yml")
    with open("faces_data/names.csv", "w") as f:
        f.write("\n".join(known_face_names))
    print("人脸数据已保存")


# 记录考勤
def record_attendance(name):
    now = datetime.now()
    date = now.strftime("%Y-%m-%d")
    time = now.strftime("%H:%M:%S")
    check_type = "上班打卡" if now.hour < 12 else "下班打卡"

    with open("attendance.csv", mode="a", newline="", encoding="utf-8") as file:
        writer = csv.writer(file)
        if file.tell() == 0:
            writer.writerow(["姓名", "日期", "时间", "考勤类型"])
        writer.writerow([name, date, time, check_type])

    print(f"{name} {check_type} 记录成功 - {date} {time}")


# 人脸注册
def register_face():
    global id_count

    name = input("请输入员工姓名: ")
    if name in known_face_names:
        print("该姓名已存在!")
        return

    cap = cv2.VideoCapture(0)
    print("请正对摄像头，按空格键拍照...")
    sample_count = 0

    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法获取摄像头画面")
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)

        for (x, y, w, h) in faces:
            cv2.rectangle(frame, (x, y), (x + w, y + h), (255, 0, 0), 2)

            # 按空格键采集样本
            if cv2.waitKey(1) & 0xFF == ord(' '):
                face_img = gray[y:y + h, x:x + w]
                face_samples.append(face_img)
                ids.append(id_count)
                sample_count += 1
                print(f"已采集 {sample_count} 张样本")

        cv2.imshow('人脸注册 - 按空格拍照 (采集20张后自动退出)', frame)

        if sample_count >= 20 or cv2.waitKey(1) & 0xFF == ord('q'):
            break

    if len(face_samples) > 0:
        recognizer.update(face_samples, np.array(ids))
        known_face_names.append(name)
        id_count += 1
        save_face_data()
        print(f"员工 {name} 注册成功!")

    cap.release()
    cv2.destroyAllWindows()


# 人脸识别考勤
def recognize_faces():
    if len(known_face_names) == 0:
        print("没有注册的员工数据，请先注册员工!")
        return

    cap = cv2.VideoCapture(0)
    print("人脸考勤中...按Q键退出")

    # 记录已打卡人员
    checked_in_today = []

    while True:
        ret, frame = cap.read()
        if not ret:
            print("无法获取摄像头画面")
            break

        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)

        for (x, y, w, h) in faces:
            face_img = gray[y:y + h, x:x + w]

            # 识别人脸
            id_, confidence = recognizer.predict(face_img)

            # 置信度低于50认为是可信识别
            if confidence < 50:
                name = known_face_names[id_]
                confidence_text = f"{round(100 - confidence)}%"

                # 检查是否已打卡
                today = datetime.now().strftime("%Y-%m-%d")
                if f"{name}_{today}" not in checked_in_today:
                    record_attendance(name)
                    checked_in_today.append(f"{name}_{today}")
            else:
                name = "Unknown"
                confidence_text = f"{round(100 - confidence)}%"

            # 绘制识别结果
            cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 2)
            cv2.putText(frame, f"{name} {confidence_text}", (x, y - 10),
                        cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 255, 0), 2)

        cv2.imshow('人脸考勤系统 - 按Q退出', frame)

        if cv2.waitKey(1) & 0xFF == ord('q'):
            break

    cap.release()
    cv2.destroyAllWindows()


# 查看考勤记录
def view_attendance():
    if not os.path.exists("attendance.csv"):
        print("暂无考勤记录")
        return

    print("\n考勤记录:")
    with open("attendance.csv", mode="r", encoding="utf-8") as file:
        reader = csv.reader(file)
        for row in reader:
            print(f"{row[0]:<10} {row[1]:<12} {row[2]:<10} {row[3]}")


# 主菜单
def main():
    load_known_faces()

    while True:
        print("\n===== 人脸考勤系统 =====")
        print("1. 员工注册")
        print("2. 人脸考勤")
        print("3. 查看考勤记录")
        print("4. 退出系统")

        choice = input("请选择功能 (1-4): ")

        if choice == '1':
            register_face()
        elif choice == '2':
            recognize_faces()
        elif choice == '3':
            view_attendance()
        elif choice == '4':
            print("退出系统...")
            break
        else:
            print("无效选择，请重新输入!")


if __name__ == "__main__":
    main()